import os

from langchain_community.chat_message_histories import SQLChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.agents import create_tool_calling_agent
from langchain.tools.retriever import create_retriever_tool

os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_PROJECT"] = "simple_robot[1.0.0]"
os.environ["LANGCHAIN_API_KEY"] = "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"
os.environ["GOOGLE_API_KEY"] = "AIzaSyBJoz7BvdFgWTBwzcu-0xWpJKfEJOR6vPM"

llm = ChatGoogleGenerativeAI(model="models/gemini-1.5-pro-latest", temperature=0.3)

connection_string = "sqlite:///F:/tmp/temp/sqlite.db"

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "You are a helpful assistant."),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{question}"),
    ]
)

chain = prompt | llm

chain_with_history = RunnableWithMessageHistory(
    chain,
    lambda session_id: SQLChatMessageHistory(
        session_id=session_id, connection_string=connection_string
    ),
    input_messages_key="question",
    history_messages_key="history",
)

config = {"configurable": {"session_id": "session_id_1"}}

## 获取聊天记录
print(chain_with_history.get_session_history("session_id_1"))

print(chain_with_history.invoke({"question": "你觉得5岁男孩适合玩什么样的玩具？"}, config=config))
